﻿#include "caculaer.h"
#include <QFile>

typedef struct {
    double m[3][3];
} Matrix3x3;
typedef struct {
    double w, x, y, z;
} Quaternion;

void quaternionToRotationMatrix(double w, double x, double y, double z, Matrix3x3 *R)
{
    R->m[0][0] = 1.0 - 2.0 * y * y - 2.0 * z * z;
    R->m[0][1] = 2.0 * x * y - 2.0 * w * z;
    R->m[0][2] = 2.0 * x * z + 2.0 * w * y;

    R->m[1][0] = 2.0 * x * y + 2.0 * w * z;
    R->m[1][1] = 1.0 - 2.0 * x * x - 2.0 * z * z;
    R->m[1][2] = 2.0 * y * z - 2.0 * w * x;

    R->m[2][0] = 2.0 * x * z - 2.0 * w * y;
    R->m[2][1] = 2.0 * y * z + 2.0 * w * x;
    R->m[2][2] = 1.0 - 2.0 * x * x - 2.0 * y * y;
}
Quaternion matrixToQuaternion(double m[3][3]) {
    Quaternion q;
    double trace = m[0][0] + m[1][1] + m[2][2];
    double s;

    if (trace > 0) {
        s = sqrt(trace + 1.0);
        q.w = s * 0.5;
        s = 0.5 / s;
        q.x = (m[2][1] - m[1][2]) * s;
        q.y = (m[0][2] - m[2][0]) * s;
        q.z = (m[1][0] - m[0][1]) * s;
    } else {
        // 查找最大的对角线元素以确定四元数的实部
        int i = 0;
        if (m[1][1] > m[0][0]) i = 1;
        if (m[2][2] > m[i][i]) i = 2;
        int j = (i + 1) % 3;
        int k = (j + 1) % 3;

        s = sqrt((m[i][i] - (m[j][j] + m[k][k])) + 1.0);
        switch (i) {
        case 0:
            q.w = (m[2][1] - m[1][2]) * 0.5 / s;
            q.x = s * 0.5;
            q.y = (m[0][1] + m[1][0]) * 0.5 / s;
            q.z = (m[0][2] + m[2][0]) * 0.5 / s;
            break;
        case 1:
            q.w = (m[0][2] - m[2][0]) * 0.5 / s;
            q.x = (m[0][1] + m[1][0]) * 0.5 / s;
            q.y = s * 0.5;
            q.z = (m[1][2] + m[2][1]) * 0.5 / s;
            break;
        case 2:
            q.w = (m[1][0] - m[0][1]) * 0.5 / s;
            q.x = (m[0][2] + m[2][0]) * 0.5 / s;
            q.y = (m[1][2] + m[2][1]) * 0.5 / s;
            q.z = s * 0.5;
            break;
        }
    }

    // 可选：四元数正规化
    double norm = sqrt(q.w * q.w + q.x * q.x + q.y * q.y + q.z * q.z);
    q.w /= norm;
    q.x /= norm;
    q.y /= norm;
    q.z /= norm;
    return q;
}
QVector3D rotationMatrixToEulerAngles(double R[3][3])
{
    double sy = sqrt(R[0][0] * R[0][0] +  R[1][0] * R[1][0]);
    bool singular = sy < 1e-6;
    double x,y,z;
    if (!singular)
    {
        x = atan2(R[2][1] , R[2][2]);
        y =atan2(-R[2][0], sy);
        z = atan2(R[1][0], R[0][0]);
    }
    else
    {
        x = atan2(-R[1][2], R[1][1]);
        y = atan2(-R[2][0], sy);
        z = 0;
    }
    return QVector3D(x, y, z);
}
void eulertomatrix(double a,double b,double c,Matrix3x3* R)
{
    double cosyaw= cos(c);
    double sinyaw= sin(c);
    double cospitch= cos(b);
    double sinpitch= sin(b);
    double cosroll= cos(a);
    double sinroll= sin(a);
    R->m[0][0] = cosyaw * cospitch;
    R->m[0][1] = cosyaw * sinpitch * sinroll - sinyaw * cosroll;
    R->m[0][2] = cosyaw * sinpitch * cosroll + sinyaw * sinroll;
    R->m[1][0] = sinyaw * cospitch;
    R->m[1][1] = sinyaw * sinpitch * sinroll + cosyaw * cosroll;
    R->m[1][2] = sinyaw * sinpitch * cosroll - cosyaw * sinroll;
    R->m[2][0] = -sinpitch;
    R->m[2][1] = cospitch * sinroll;
    R->m[2][2] = cospitch * cosroll;
}

Caculaer::Caculaer()
{

}

Caculaer::~Caculaer()
{

}

void Caculaer::clear()
{
    fmap.clear();
}

void Caculaer::setdatasouce(QString path)
{
    imagefilepath=path.toStdString();
    std::cout<<"image file path == "<<imagefilepath<<std::endl;

}

bool Caculaer::execute(const QMap<int, Freme>& source)
{

    if(source.size()==0)
    {
        LOG_RED("SOMETHING WRONG WHEN LOADING IMAGES OR ROBOT POSES");
        return false;
    }
    fmap = source;
    if(!loadImages())
    {
        LOG_ERROR("LOAD DATA FAILED!");
        return false;
    }
    if(!loadPoses())
    {
        LOG_ERROR("LOAD POS FAILED!");
        return false;
    }
    if (!loadCamParams()) {
        LOG_RED("UNABLE TO LOAD CAMERA PARAMETERS");
        return false;
    }
    if (m_input_imgs.size() == 0 || m_robot_poses.size() == 0) {
        LOG_RED("SOMETHING WRONG WHEN LOADING IMAGES OR ROBOT POSES");
        return false;
    }
    if (m_input_imgs.size() != m_robot_poses.size()) {
        LOG_RED("NUMBER OF IMAGES NEEDS TO AGREE WITH NUMBER OF POSES");
        return false;
    }

    std::vector<cv::Mat> R_target2cam,t_target2cam;
    cv::Size board_size(cornerColum,cornerRow);
    Pattern p = Pattern::CHESSBOARD;
    float board_dim=18;
    for(int i=0;i<m_input_imgs.size();++i)
    {
        std::vector<cv::Point2f> centers_output;
        cv::Mat rvec, tvec;
        try{
            if (!findPattern(m_input_imgs[i], board_size, p, m_img_names[i], centers_output)) {
                LOG_RED("CANNOT FIND PATTERN FOR IMAGE %i", i);
                return false;
            }
            if (!findBoardPose(board_size, board_dim, centers_output, m_intrinsics, m_distortion, rvec, tvec)) {
                LOG_RED("CANNOT CALCULATE BOARD POSE");
                return false;
            }
            cv::Mat rot;
            cv::Rodrigues(rvec, rot);
            R_target2cam.push_back(rot);
            t_target2cam.push_back(tvec);
        } catch (cv::Exception& e) {
            LOG_RED(e.what());
            return false;
        } catch (std::exception& e) {
            LOG_RED(e.what());
            return false;
        }
    }
    std::vector<cv::Mat> R_gripper2base, t_gripper2base;
    for (int i = 0; i < m_robot_poses.size(); i++) {
        Eigen::Isometry3f pose = m_robot_poses[i].getHomoPose();
        Eigen::Matrix3f rot_mat = pose.rotation();
        cv::Mat rot = (cv::Mat_<double>(3, 3) << rot_mat(0, 0), rot_mat(0, 1), rot_mat(0, 2),
                       rot_mat(1, 0), rot_mat(1, 1), rot_mat(1, 2),
                       rot_mat(2, 0), rot_mat(2, 1), rot_mat(2, 2));
        Eigen::Vector3f trans_mat = pose.translation();
        cv::Mat trans = (cv::Mat_<double>(3, 1) << trans_mat(0), trans_mat(1), trans_mat(2));
        R_gripper2base.push_back(rot);
        t_gripper2base.push_back(trans);
    }
    try{
        cv::calibrateHandEye(
            R_gripper2base,
            t_gripper2base,
            R_target2cam,
            t_target2cam,
            m_R_camera2gripper,
            m_t_camera2gripper,
            cv::CALIB_HAND_EYE_TSAI
            );

        // perform calibration verfication here
        auto RT2Homo = [](cv::Mat R, cv::Mat t)->cv::Mat{
            cv::Mat homo = (cv::Mat_<double>(4, 4) << R.at<double>(0, 0), R.at<double>(0, 1), R.at<double>(0, 2), t.at<double>(0, 0),
                            R.at<double>(1, 0), R.at<double>(1, 1), R.at<double>(1, 2), t.at<double>(1, 0),
                            R.at<double>(2, 0), R.at<double>(2, 1), R.at<double>(2, 2), t.at<double>(2, 0),
                            0.0               , 0.0               , 0.0               , 1.0);
            return homo;
        };
        m_target2base.clear();
        cv::Mat H_c2g = RT2Homo(m_R_camera2gripper, m_t_camera2gripper);
        m_Homo_camera2gripper = H_c2g;
        for (int i = 0; i < R_gripper2base.size(); i++) {
            cv::Mat H_g2b = RT2Homo(R_gripper2base[i], t_gripper2base[i]);
            cv::Mat H_t2c = RT2Homo(R_target2cam[i], t_target2cam[i]);
            cv::Mat H_t2b = H_g2b * H_c2g * H_t2c;
            m_target2base.push_back(H_t2b);
        }
        std::cout<<"rotation result == \n"<<m_R_camera2gripper<<std::endl;
        std::cout<<"translate reslut == \n"<<m_t_camera2gripper<<std::endl;

        std::cout<<m_Homo_camera2gripper<<std::endl;

    /*
    time(&m_rawtime);
    m_time_of_calib = asctime(localtime(&m_rawtime));

    if (save_result == true) {
            cv::FileStorage fs_eih(m_result_dir + "/" + "eih_result.yml", cv::FileStorage::WRITE);
            fs_eih << "CalibrationTime" << m_time_of_calib;
            fs_eih << "DistanceUnit" << dist_unit;
            fs_eih << "Cam2Gripper_Rot" << m_R_camera2gripper;
            fs_eih << "Cam2Gripper_Trans" << m_t_camera2gripper;
            fs_eih << "Cam2Gripper_Homo" << H_c2g;
            fs_eih.release();

            cv::FileStorage fs_t2b(m_result_dir + "/" + "t2b_result.yml", cv::FileStorage::WRITE);
            fs_t2b << "CalibrationTime" << m_time_of_calib;
            fs_t2b << "DistanceUnit" << dist_unit;
            for (int i = 0; i < m_target2base.size(); i++) {
                std::string pose_name = "pose" + std::to_string(i + 1);
                fs_t2b << pose_name << m_target2base[i];
            }
            fs_t2b.release();
    }
    m_calib_performed = true;
    return true;
*/
    } catch (cv::Exception& e) {
    LOG_RED(e.what());
    return false;
    } catch (std::exception& e) {
    LOG_RED(e.what());
    return false;
    }
    return true;
}

QMatrix4x4 Caculaer::getmatrixresult()
{
    QMatrix4x4 mat;
    for(int i=0;i<4;++i)
    {
        for(int j=0;j<4;++j)
        {
            mat(i,j) = m_Homo_camera2gripper.at<double>(i,j);
        }
    }
    return mat;
}

QVector<double> Caculaer::geteulerreslut()
{
    QVector<double> pos;
    qDebug()<<"m_t_camera2gripper.rows == "<<m_t_camera2gripper.rows;
    qDebug()<<"m_t_camera2gripper.cols == "<<m_t_camera2gripper.cols;
    pos<<m_t_camera2gripper.at<double>(0, 0)<<m_t_camera2gripper.at<double>(1, 0)<<m_t_camera2gripper.at<double>(2, 0);

    Matrix3x3 mat;
    mat.m[0][0] = m_R_camera2gripper.at<double>(0,0);mat.m[0][1] = m_R_camera2gripper.at<double>(0,1);mat.m[0][2] = m_R_camera2gripper.at<double>(0,2);
    mat.m[1][0] = m_R_camera2gripper.at<double>(1,0);mat.m[1][1] = m_R_camera2gripper.at<double>(1,1);mat.m[1][2] = m_R_camera2gripper.at<double>(1,2);
    mat.m[2][0] = m_R_camera2gripper.at<double>(2,0);mat.m[2][1] = m_R_camera2gripper.at<double>(2,1);mat.m[2][2] = m_R_camera2gripper.at<double>(2,2);
    QVector3D eul = rotationMatrixToEulerAngles(mat.m);
    pos<<eul.x()<<eul.y()<<eul.z();
    return pos;
}

QVector<double> Caculaer::getfourelementresult()
{
    QVector<double> pos;
    Matrix3x3 mat;
    mat.m[0][0] = m_R_camera2gripper.at<double>(0,0);mat.m[0][1] = m_R_camera2gripper.at<double>(0,1);mat.m[0][2] = m_R_camera2gripper.at<double>(0,2);
    mat.m[1][0] = m_R_camera2gripper.at<double>(1,0);mat.m[1][1] = m_R_camera2gripper.at<double>(1,1);mat.m[1][2] = m_R_camera2gripper.at<double>(1,2);
    mat.m[2][0] = m_R_camera2gripper.at<double>(2,0);mat.m[2][1] = m_R_camera2gripper.at<double>(2,1);mat.m[2][2] = m_R_camera2gripper.at<double>(2,2);

    Quaternion quant;
    quant = matrixToQuaternion(mat.m);
    pos<<m_t_camera2gripper.at<double>(0, 0)<<m_t_camera2gripper.at<double>(1, 0)<<m_t_camera2gripper.at<double>(2, 0);
    pos<<quant.w<<quant.x<<quant.y<<quant.z;
    return pos;
}

bool Caculaer::loadImages()
{
    for(int i=0;i< fmap.keys().size();++i)
    {
        int index= fmap.keys().at(i);
        try{
            std::string imgname=fmap[i].imagepath.toStdString();
            //imgname=imagefilepath+"/"+imgname;
            if(!(QString::fromStdString(imgname).endsWith(".png")||QString::fromStdString(imgname).endsWith(".jpg")))
            {
                LOG_RED("ONLY SUPPORT PNG AND JPG IMAGES. OR THE DIRECTORY DOES NOT CONTIAN IMAGES");
                return false;
            }
            m_img_names.push_back(imgname);
        }catch(cv::Exception& e)
        {
            LOG_RED(e.what());
            return false;
        }
    }
    if (m_img_names.size() < 4) {
        LOG_RED("NEED AT LEAST 4 IMAGES IN THE DIRECTORY");
        return false;
    }

    for (int i = 0; i < m_img_names.size(); i++) {
        try{
            std::string dir = imagefilepath + "/" + m_img_names[i];
            QFile file(QString::fromStdString(dir));
            if(!file.exists())
            {
                LOG_RED("CAN NOT FIND FILE %S",dir.c_str());
                return false;
            }
            cv::Mat img = cv::imread(dir);
            m_input_imgs.push_back(img);
        } catch (cv::Exception& e) {
            LOG_RED("%s", e.what());
            return false;
        }
    }

    LOG_GREEN("%d IMAGE LOADED SUCCESSFULLY",m_input_imgs.size());
    cv::Size img_size= m_input_imgs[0].size();

    for (auto img : m_input_imgs) {
        if (img.size() != img_size) {
            LOG_ERROR("IMAGES ARE OF DIFFERENT SIZE, CANNOT PERFORM CALIBRATION");
            m_input_imgs.clear();
            return false;
        }
    }
    return true;
}

bool Caculaer::loadPoses()
{
    m_robot_poses.clear();
    for(int i=0;i<fmap.keys().size();++i)
    {
        int index=fmap.keys().at(i);
        QVector<double> pos_vals = fmap[index].pos;
        RobotPose rp;
        rp.updateTrans(pos_vals[0], pos_vals[1], pos_vals[2]);
        rp.updateRotVec(pos_vals[3], pos_vals[4], pos_vals[5]);
        m_robot_poses.push_back(rp);
    }
    LOG_GREEN("%d ROBOT POSE LOADED SUCCESSFULLY",m_robot_poses.size());
    return true;
}

bool Caculaer::loadCamParams()
{
    QString path=QString::fromStdString(imagefilepath)+"/intrinsics.yml";
    QFile file(path);
    if(!file.exists())
    {
        LOG_RED("DIRECTORY DOES NOT EXIST");
        return false;
    }

    cv::FileStorage fs(path.toStdString(), cv::FileStorage::READ);
    if (!fs.isOpened()) {
        LOG_RED("UNABLE TO OPEN YAML POSE FILE");
        return false;
    }
    try {
        fs["Intrinsic"] >> m_intrinsics;
        fs["Distortion"] >> m_distortion;
        fs["ImageSize"] >> m_expect_size;
        fs.release();
        LOG_GREEN("CAMERA PARAMETERS LOADED SUCCESSFULLY");
        return true;
    } catch (cv::Exception& e) {
        LOG_RED(e.what());
        return false;
    } catch (std::exception& e) {
        LOG_RED(e.what());
        return false;
    }
    return true;
}

bool Caculaer::findPattern(const cv::Mat &input, const cv::Size &board_size, const Pattern &p, const std::string &img_name, std::vector<cv::Point2f> &centers_output, bool draw_result)
{
    centers_output.clear();
    switch(p) {
    case Pattern::CHESSBOARD : {
        cv::Mat track_result;
        bool pattern_found = cv::findChessboardCorners(input, board_size, centers_output);
        if (!pattern_found) {
            LOG_RED("UNABLE TO FIND CHESSBOARD PATTERN");
            cv::imshow(img_name, input);
            cv::waitKey(0);
            cv::destroyWindow(img_name);
            return false;
        } else {
            // refine to subpixel accuracy
            cv::Mat gray;
            cv::cvtColor(input, gray, cv::COLOR_BGR2GRAY);
            cv::cornerSubPix(
                gray,
                centers_output,
                cv::Size(11, 11),
                cv::Size(-1, -1),
                cv::TermCriteria(cv::TermCriteria::EPS + cv::TermCriteria::COUNT, 30, 0.01)
                );
            if (draw_result == true) {
                cv::Mat display = input.clone();
                cv::drawChessboardCorners(display, board_size, centers_output, pattern_found);
                cv::imshow("Found corners", display);
                cv::waitKey(0);
                cv::destroyWindow("Found corners");
            }
            return true;
        }

        break;
    }
    case Pattern::SYMCIRCLE : {
        cv::Mat track_result;
        bool pattern_found = cv::findCirclesGrid(input, board_size, track_result);
        if (!pattern_found) {
            LOG_RED("UNABLE TO FIND CIRCLE GRID PATTERN");
            cv::imshow("Track Failed", input);
            cv::waitKey(0);
            cv::destroyWindow("Track Failed");
            return false;
        } else {
            cv::Mat display = input.clone();
            for (int i = 0; i < track_result.rows; i++) {
                centers_output.push_back(track_result.at<cv::Point2f>(0, i));
            }
            if (draw_result == true) {
                cv::drawChessboardCorners(display, board_size, centers_output, pattern_found);
                cv::imshow("Found circles", display);
                cv::waitKey(0);
                cv::destroyWindow("Found circles");
            }
            return true;
        }

        break;
    }
    case Pattern::ASYMCIRCLE : {
        LOG_RED("ASYMMETRIC CIRCLES COMING SOON");
        return false;
        break;
    }
    }
    return true;
}
bool Caculaer::findBoardPose(const cv::Size &board_size, const float &board_dim, const std::vector<cv::Point2f> &tracked_centers, const cv::Mat &camera_matrix, const cv::Mat &dist_coeffs, cv::Mat &rvec, cv::Mat &tvec)
{
    std::vector<cv::Point3f> object_points;
    for (int i = 0; i < board_size.height; i++) {
        for (int j = 0; j < board_size.width; j++) {
            object_points.emplace_back(i*board_dim, j*board_dim, 0);
        }
    }

    /* Calculate translation and rotation with solvePnP */
    cv::Mat rvector = cv::Mat::zeros(3, 1, CV_64F);
    cv::Mat tvector = cv::Mat::zeros(3, 1, CV_64F);
    cv::Mat tracked_centers_map = cv::Mat(1, tracked_centers.size(), CV_32FC2);
    for (int i = 0; i < tracked_centers.size(); i++) {
        tracked_centers_map.at<cv::Point2f>(0, i) = tracked_centers[i];
    }
    try {
        if (!cv::solvePnP(object_points, tracked_centers_map, camera_matrix, dist_coeffs, rvector, tvector)) {
            return false;
        };
        rvec = rvector.clone();
        tvec = tvector.clone();
        return true;
    } catch (cv::Exception& e) {
        LOG_RED(e.what());
        return false;
    } catch (std::exception& e) {
        LOG_RED(e.what());
        return false;
    }
}

/*
bool Caculaer::findPattern(Freme& frem,bool draw_result)//const cv::Mat &input, const cv::Size board_size, const Pattern &P, const std::string &img_name, std::vector<cv::Point2f> &center_output, bool draw_result)
{
    frem.center.clear();
    cv::Mat tracl_result;
    cv::Size board_size(cornerRow,cornerColum);
    bool pattten_found = cv::findChessboardCorners(frem.image,board_size,frem.center);
    std::string img_name = frem.imagepath.toStdString();
    if(!pattten_found)
    {
        LOG_RED("UNABLE TO FIND CHESSBOARD PATTERN");
        cv::imshow(img_name, frem.image);
        cv::waitKey(0);
        cv::destroyWindow(img_name);
        return false;
    }
    else {
        cv::Mat gray;
        cv::cvtColor(frem.image, gray, cv::COLOR_BGR2GRAY);
        cv::cornerSubPix(
            gray,
            frem.center,
            cv::Size(11, 11),
            cv::Size(-1, -1),
            cv::TermCriteria(cv::TermCriteria::EPS + cv::TermCriteria::COUNT, 30, 0.01)
            );
        if (draw_result == true) {
            cv::Mat display = frem.image.clone();
            cv::drawChessboardCorners(display, board_size, frem.center, pattten_found);
            cv::imshow("Found corners", display);
            cv::waitKey(0);
            cv::destroyWindow("Found corners");
        }
        return true;
    }

    return pattten_found;
}
*/
void Caculaer::RobotPose::updateTrans(double x, double y, double z)
{
    this->x=x;
    this->y=y;
    this->z=z;
    this->is_trans_set=true;
}

void Caculaer::RobotPose::updateQuat(double qw, double qx, double qy, double qz)
{

}

void Caculaer::RobotPose::updateRotVec(double rx, double ry, double rz)
{
    Eigen::Vector3f rot(rx,ry,rz);
    this->rot_vec.angle()=rot.norm();
    this->rot_vec.axis()=rot.normalized();
    this->quat=this->rot_vec;
    this->rot_mat=this->quat;
    this->is_rot_set=true;

}

void Caculaer::RobotPose::updateEulerAngle(double r, double p, double y)
{

}

std::vector<double> Caculaer::RobotPose::getPose(std::string rot_rep)
{
    return std::vector<double>();
}

std::vector<double> Caculaer::RobotPose::getPoseRPY(std::string rot_order)
{
    return std::vector<double>();
}

Eigen::Isometry3f Caculaer::RobotPose::getHomoPose()
{
    Eigen::Vector3f trans(this->x,this->y,this->z);
    Eigen::Isometry3f pose= Eigen::Isometry3f::Identity();
    pose.pretranslate(trans);
    pose.rotate(this->quat);
    return pose;
}
